'''
Author: Zhuangyu
Date: 2022-10-21 15:32:15
LastEditTime: 2022-10-21 16:29:11
'''
from __future__ import print_function, division
import time
#Importing relevant modules and packages
import numpy as np
import matplotlib.pyplot as plt
import random
from Parameters import *
from scipy import integrate, linalg, interpolate
from tqdm import tqdm
from casadi import *
import mpctools as mpc
import casadi
import scipy.io as sio  
import van_Genuchten as vg
from RichardsModel_1D import *
from prejacobian import F1,F2,F3,Fy

#Space and time parameters
Nr,Nt,Nz,dr,dt,dz,Np,Na,Np,N_obs,Hz=circular_parameters()
DeltaT,Nsim=time_parameters()
p=Loam() #Soil parameters


size=Nsim

def ode_sensitivity(x,u,ui):
    
    Sxx = np.zeros([size+1,Np,Np])
    S1=np.eye(Nz)
    S2=np.eye(Na)
    S11=np.hstack((S1,S2))
    S22=np.hstack((np.zeros((Nz,Nz)),np.eye(Nz)))
    Sxx[0,:,:] = np.concatenate((S11,S22),axis=0)
    Syx = np.zeros([size+1,N_obs,Np])
    
    dFdx = np.zeros([size,Nz,Nz])
    dFdu = np.zeros([size,Na,Na])
    dF=np.zeros([size,Np,Np])
    dHdx = np.zeros([size+1,N_obs,Nz])
    dH=np.zeros([size+1,N_obs,Np])
        
    
    for i in range (size):
        dFdx[i,:,:]=F1(x[i,:],u[i,:],ui[i,:],DeltaT)
        dFdu[i,:,:]=F3(x[i,:],u[i,:],ui[i,:],DeltaT)
        dF1=np.hstack((dFdx[i,:,:],dFdu[i,:,:]))
        dF2=np.hstack((np.zeros((Nz,Nz)),np.eye(Nz)))
        dF[i,:,:]=np.concatenate((dF1,dF2),axis=0)
        dHdx[i,:,:]=Fy(x[i,:],u[i,:],ui[i,:],DeltaT)
        dH[i,:,:]=np.hstack((dHdx[i,:,:],np.zeros((N_obs,Na))))
                 
        Sxx[i+1,:,:]=mtimes(dF[i,:,:],Sxx[i,:,:])
        
        Syx[i,:,:]=mtimes(dH[i,:,:],Sxx[i,:,:])
    
    for i in range (size+1):
        if i==0:
            Sall = Syx[i,:,:]
        else:
            Sall = np.concatenate((Sall,Syx[i,:,:]),axis=0)
    
    return Sall
        